import cv2
import numpy as np
import os
import sys


def write_line(data,yuv_file_name):
	with open('output.txt', 'a') as file:
		file.write("input:"+yuv_file_name+"       ")
		file.write(data)
		file.write('\n')

def read_yuv420_frame(width, height, file, frame_index):
    """
    从 YUV420 文件中读取指定帧的 Y、U、V 分量。
    """
    frame_size = width * height * 3 // 2
    file.seek(frame_size * frame_index)
    yuv_data = file.read(frame_size)
    if len(yuv_data) != frame_size:
        return None  # 读取不到足够数据时，返回 None

    y = np.frombuffer(yuv_data[:width * height], dtype=np.uint8).reshape((height, width))
    u = np.frombuffer(yuv_data[width * height:width * height + (width * height) // 4], dtype=np.uint8).reshape((height // 2, width // 2))
    v = np.frombuffer(yuv_data[width * height + (width * height) // 4:], dtype=np.uint8).reshape((height // 2, width // 2))

    return y, u, v

def yuv420_to_bgr(y, u, v, width, height):
    """
    将 YUV420 图像转换为 BGR 图像。
    """
    yuv_image = np.zeros((height + height // 2, width), dtype=np.uint8)
    yuv_image[:height, :] = y
    yuv_image[height:, :width // 2] = u
    yuv_image[height:, width // 2:] = v

    bgr_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_YV12)
    return bgr_image

def calculate_psnr(image1, image2):
    """
    计算两个图像的 PSNR 值。
    """
    return cv2.PSNR(image1, image2)

def process_yuv_files(width, height, file1_path, file2_path, frame_count):
    """
    对两个 YUV 文件的多帧计算 PSNR。
    """
    psnr_values = []

    with open(file1_path, 'rb') as file1, open(file2_path, 'rb') as file2:
        for frame_index in range(frame_count):
            # 读取两帧
            y1, u1, v1 = read_yuv420_frame(width, height, file1, frame_index)
            y2, u2, v2 = read_yuv420_frame(width, height, file2, frame_index)

            if y1 is None or y2 is None:
                print(f"End of file reached at frame {frame_index}")
                break

            # 转换为 BGR 图像
            bgr1 = yuv420_to_bgr(y1, u1, v1, width, height)
            bgr2 = yuv420_to_bgr(y2, u2, v2, width, height)

            # 计算 PSNR
            psnr_value = calculate_psnr(bgr1, bgr2)
            psnr_values.append(psnr_value)
            #print(f"Frame {frame_index + 1}: PSNR = {psnr_value:.2f} dB")

    # 计算平均 PSNR
    avg_psnr = sum(psnr_values) / len(psnr_values) if psnr_values else 0
    print(f"Average PSNR: {avg_psnr:.2f} dB")

    write_line(f"Average PSNR: {avg_psnr:.2f} dB",file1_path)
    return psnr_values, avg_psnr


if __name__ == "__main__":
    file1_path = './temp.yuv'
	# 执行计算
    cmd = "/home/lzh/lzh/code/publish_new/seirios_dev/ffmpeg/ffmpeg -y -i " +sys.argv[1] + " -vcodec rawvideo -s 1920x1080 -pix_fmt yuv420p  -vframes 300 " + file1_path
    os.system(cmd)
	# 图像宽度和高度
    width = 1920
    height = 1080
	# 文件路径
    file2_path = sys.argv[2]

	# 总帧数（根据实际情况调整，或者根据文件大小推断）
    frame_count = 300

    psnr_values, avg_psnr = process_yuv_files(width, height, file1_path, file2_path, frame_count)

